# 子树的大小(好题)
from leetcode import test_function


def get_node_num(n, m, k):
    """我的方法, 对提供的三个测试用例都能提供正确答案, 但是提交代码的正确率为0%"""
    target_depth, curr_id = 0, 1  # 根节点为第0层

    # 获得目标节点的深度
    while curr_id < k:
        target_depth += 1
        curr_id += m * target_depth

    # 获得目标节点的位置(左数第几个)
    layer_id = 1
    if k != 1:
        layer_id = k - curr_id + pow(m, target_depth)  # 该层的第几个

    # 计算整个树的最终深度
    tree_depth = target_depth
    while curr_id < n:
        tree_depth += 1
        curr_id += pow(m, tree_depth)

    if target_depth == tree_depth:
        return 1

    # 计算子树的节点数
    res = 0
    for i in range(tree_depth - target_depth):
        res += pow(m, i)

    curr_id -= pow(m, tree_depth)
    curr_id += (layer_id - 1) * pow(m, tree_depth - target_depth)
    remain = min(n - curr_id, pow(m, tree_depth - target_depth))
    if remain > 0:
        res += remain

    return res


def get_node_num_opt(n, m, k):
    """deepseek实现的函数, 通过dfs的方式来实现, 超时/通过 = 5/5"""
    if k > n:
        return 0
    count = 1  # 当前节点自身
    # 计算第一个子节点的编号：m*(k-1) + 2
    first_child = m * (k - 1) + 2
    # 遍历所有m个子节点
    for i in range(m):
        child = first_child + i
        if child > n:
            break
        count += get_node_num_opt(n, m, child)
    return count


def get_node_num_official(n, m, k):
    """官方题解(非常好)
    定义了两个指针, 表示子树的最下层的最左端和最右端的两个节点
    两个公式非常好:
    节点编号为k的左子节点编号 = k * m - m + 2
    节点编号为k的右子节点编号 = k * m + 1
    """
    l, r = k, k  # 表示子节点k的做左右子端点
    res = 1  # 记录结果
    layer_num = 1  # 记录一层的k节点子树的子节点数量
    while r * m + 1 <= n:  # 如果k子节点子树的最右端的节点编号>n, 说明已经是该树的倒数第二层, 下一层必定不是满叶子节点
        layer_num *= m  # 计算该层的所有子节点数量
        l = l * m - m + 2  # 更新左端点
        r = r * m + 1  # 更新右端点
        res += layer_num
    l = l * m - m + 2  # 最后一层该子树的最左端节点的位置
    res += max(0, n - l + 1)  # 如果最后一层的叶子节点没能铺到该子树下, 那n-l+1有可能是负数, 若为负数, 该层属于k子树就是0个节点
    return res


if __name__ == '__main__':
    # n = int(input())
    # trees = []
    # for i in range(n):
    #     trees.append(tuple(map(int, input().split())))
    #
    # for i in range(n):
    #     print(get_subtree_node_num(trees[i][0], trees[i][1], trees[i][2]))

    inp = [{"n": 1, "m": 2, "k": 1},
           {"n": 11, "m": 3, "k": 4},
           {"n": 74, "m": 5, "k": 3},
           {"n": 554531897, "m": 2, "k": 161934367}, ]
    out = [1, 2, 24, 3]

    test_function(get_node_num, inp, out)
    test_function(get_node_num_opt, inp, out)
    test_function(get_node_num_official, inp, out)
